import os
from scipy.misc import imsave
import numpy as np
import pickle

CURRENT_DIR=os.path.abspath(os.path.join(os.getcwd(), ".."))
DATA_PATH=os.path.join(CURRENT_DIR,"data","cifar-10-python")
IMAGES_PATH=os.path.join(CURRENT_DIR,"data","img")

def init_prepare():
    if not os.path.exists(IMAGES_PATH):
        os.makedirs(os.path.join(IMAGES_PATH,"train"))
        os.makedirs(os.path.join(IMAGES_PATH,"test"))
    else:
        lists = os.listdir(os.path.join(IMAGES_PATH,"train"))
        for i in lists:
            os.remove(os.path.join(IMAGES_PATH,"train", i))

# 解压缩，返回解压后的字典
def unpickle(file):
    fo = open(file, 'rb')
    dict = pickle.load(fo, encoding='latin1')
    fo.close()
    return dict

# 生成训练集图片，如果需要png格式，只需要改图片后缀名即可。
def generate_train_image():
    for j in range(1, 6):
        data_name = os.path.join(DATA_PATH, ("data_batch_" + str(j)))
        train_data_dict = unpickle(data_name)
        print(data_name + " is loading...")
        for i in range(0, 10000):
            img = np.reshape(train_data_dict['data'][i], (3, 32, 32))
            img = img.transpose(1, 2, 0)  # 读取image
            img_name = os.path.join(IMAGES_PATH,"train",(str(train_data_dict['labels'][i]) + '_' + str(i + (j - 1)*10000) + '.jpg'))
            imsave(img_name, img)
        print(data_name + " loaded.")

# 生成测试集图片
def generate_test_image():
    data_name = os.path.join(DATA_PATH, "test_batch")
    test_data_dict = unpickle(data_name)
    print(data_name + " is loading...")
    for i in range(0, 10000):
        img = np.reshape(test_data_dict['data'][i], (3, 32, 32))
        img = img.transpose(1, 2, 0)
        img_name = os.path.join(IMAGES_PATH,"test",(str(test_data_dict['labels'][i]) + '_' + str(i) + '.jpg'))
        imsave(img_name, img)
    print(data_name+" loaded.")

def test():
    init_prepare()
    generate_train_image()
    generate_test_image()

if __name__ == '__main__':
    test()